#include <arch/arch.h>
#include <stdarg.h>
#include <drivers/kernel_logger.h>
#include <irq/irq_manager.h>
#include <mm/mm.h>


// 通用寄存器打印
void dump_registers(struct pt_regs *regs)
{
    // 通用寄存器部分
    printk("ra: 0x%016lx  sp: 0x%016lx\n", regs->ra, regs->sp);
    printk("gp: 0x%016lx  tp: 0x%016lx\n", regs->gp, regs->tp);
    printk("t0: 0x%016lx  t1: 0x%016lx\n", regs->t0, regs->t1);
    printk("t2: 0x%016lx  s0: 0x%016lx\n", regs->t2, regs->s0);
    printk("s1: 0x%016lx  a0: 0x%016lx\n", regs->s1, regs->a0);
    printk("a1: 0x%016lx  a2: 0x%016lx\n", regs->a1, regs->a2);
    printk("a3: 0x%016lx  a4: 0x%016lx\n", regs->a3, regs->a4);
    printk("a5: 0x%016lx  a6: 0x%016lx\n", regs->a5, regs->a6);
    printk("a7: 0x%016lx  s2: 0x%016lx\n", regs->a7, regs->s2);
    printk("s3: 0x%016lx  s4: 0x%016lx\n", regs->s3, regs->s4);
    printk("s5: 0x%016lx  s6: 0x%016lx\n", regs->s5, regs->s6);
    printk("s7: 0x%016lx  s8: 0x%016lx\n", regs->s7, regs->s8);
    printk("s9: 0x%016lx  s10: 0x%016lx\n", regs->s9, regs->s10);
    printk("s11: 0x%016lx  t3: 0x%016lx\n", regs->s11, regs->t3);
    printk("t4: 0x%016lx  t5: 0x%016lx\n", regs->t4, regs->t5);
    printk("t6: 0x%016lx\n", regs->t6); // 最后一个单独一行

    // CSR 寄存器部分
    printk("epc: 0x%016lx  sstatus: 0x%016lx\n", regs->epc, regs->sstatus);
    printk("stval: 0x%016lx\n", regs->stval);
}

void handle_trap_c(struct pt_regs *regs){
    uint64_t is_interrupt = csr_read(scause) & (1UL << 63);
    uint64_t cause_code = csr_read(scause) & 0x7FFFFFFFFFFFFFFF;
    
    if(is_interrupt){
        handle_interrupt_c(regs, cause_code);
    }else{
        printk("Exception occurred: \n");

        dump_registers(regs);

        handle_exception_c(regs, cause_code);
    }
}

void handle_exception_c(struct pt_regs *regs, uint64_t cause)
{
    switch (cause) {
        case 0:
            printk("[EXC] Instruction address misaligned at 0x%lx\n", regs->epc);
            break;
        case 1:
            printk("[EXC] Instruction access fault at 0x%lx\n", regs->epc);
            break;
        case 2:
            printk("[EXC] Illegal instruction at 0x%lx\n", regs->epc);
            break;
        case 3:
            printk("[EXC] Breakpoint at 0x%lx\n", regs->epc);
            regs->epc += 4; // 跳过 ebreak
            return;
        case 4:
            printk("[EXC] Load address misaligned at 0x%lx\n", regs->stval);
            break;
        case 5:
            printk("[EXC] Load access fault at 0x%lx\n", regs->stval);
            break;
        case 6:
            printk("[EXC] Store/AMO address misaligned at 0x%lx\n", regs->stval);
            break;
        case 7:
            printk("[EXC] Store/AMO access fault at 0x%lx\n",regs->stval);
            break;
        case 8:
        printk("[EXC] Environment call from U-mode (syscall) sepc=0x%lx\n", regs->epc);

        /* 1) 判定指令长度（默认 4）——安全地短期开 SUM 访问用户内存 */
        uint64_t instr_len = 4;
        uint16_t maybe_half = 0;
    
        /* 临时打开 SUM 以便 S-mode 访问 U-mode 内存（只在需要时） */
        uint64_t sstatus_bak;
        __asm__ volatile("csrr %0, sstatus" : "=r"(sstatus_bak));
        __asm__ volatile("csrs sstatus, %0" :: "r"(1UL << 18) : "memory");
    
        /* 读取用户指令的低16位来判断是否为 compressed (2-byte) */
        /* 防御：把读取包在 try/guard 中很困难，暂假设映射有效（或你可以先用 translate_address 验证） */
        maybe_half = *(uint16_t *)(uintptr_t)(regs->epc);
        instr_len = ((maybe_half & 0x3) != 0x3) ? 2 : 4;
    
        /* 恢复 sstatus（清除 SUM） */
        __asm__ volatile("csrw sstatus, %0" :: "r"(sstatus_bak) : "memory");
    
        /* 2) 先把 sepc 前进（写入 regs 内存）——确保返回点是下一条指令 */
        uint64_t old_epc = regs->epc;
        regs->epc = old_epc + instr_len;
    
        /* 防御性写回：确保保存到 trap frame 的内存中（regs 指向那块内存，但我们显式写一次） */
        uint64_t *sepc_ptr = (uint64_t *)((char *)regs + offsetof(struct pt_regs, epc));
        *sepc_ptr = regs->epc;
    
        printk("[syscall] before call: old_epc=0x%lx new_epc=0x%lx a7=0x%lx\n",
               old_epc, regs->epc, regs->a7);
    
        /* 3) 调用 syscall 分发器：它应返回 uint64_t，并**不要**去改 regs->epc（除非特殊 syscall） */
        uint64_t ret = handle_syscall(regs);
        regs->a0 = ret;
    
        /* 再次防御性写回 sepc（以防handle_syscall意外改写） */
        *sepc_ptr = regs->epc;
    
        printk("[syscall] after call: ret=0x%lx new_epc=0x%lx\n", regs->a0, regs->epc);
    
        /* 断言：如果 handle_syscall 非预期修改了 sepc，记录并继续（便于调试） */
        if (regs->epc == old_epc) {
            printk("[BUG] syscall handler left sepc unchanged (loop risk). old_epc=0x%lx\n", old_epc);
        }
    
        return;
        case 9:
            printk("[EXC] Environment call from S-mode\n");
            regs->epc += 4;
            return;
        case 12:
            printk("[EXC] Instruction page fault at VA=0x%lx\n", regs->stval);
            break;
        case 13:
            printk("[EXC] Load page fault at VA=0x%lx\n", regs->stval);
            break;
        case 15:
            printk("[EXC] Store/AMO page fault at VA=0x%lx\n", regs->stval);
            break;
        default:
            printk("[EXC] Unknown exception %lu at 0x%lx stval=0x%lx\n",
                   cause, regs->epc, regs->stval);
            break;
    }

    dump_registers(regs);
    while (1)
        arch_pause();
}

extern void riscv64_timer_handler(struct pt_regs *regs);

extern void do_irq(struct pt_regs *regs, uint64_t irq_num);

bool can_schedule = true;

void handle_interrupt_c(struct pt_regs *regs, uint64_t cause)
{
    switch (cause) {
        case 5:

            riscv64_timer_handler(regs);
            if(can_schedule) {
                arch_task_switch_to(regs, current_task,
                    task_search(TASK_READY, current_cpu_id));
            }
            break;
        
        default:
            do_irq(regs, cause);
            break;
    }
}

extern int init_trap_vector();

// 安全的初始化函数
int trap_init(void) {
    // 调用汇编初始化函数
    int result = init_trap_vector();
    if (result != 0) {
        printk("Failed to initialize trap vector: %d\n", result);
        return result;
    }

    // 使能机器模式中断
    uint64_t sstatus;
    asm volatile("csrr %0, sstatus" : "=r"(sstatus));
    sstatus |= (1 << 3); // 设置SIE位
    asm volatile("csrw sstatus, %0" ::"r"(sstatus));

    return 0;
}